#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

#define N 10
#define FINISH_TAG 2*N

int main(int argc, char* argv[]){
	int i, j, rank, size;
	
	int vector[N], row[N];	

	MPI_Status status;
	
	
	
	MPI_Init(&argc, &argv);
	MPI_Comm_rank(MPI_COMM_WORLD,&rank);
	MPI_Comm_size(MPI_COMM_WORLD,&size);
	
	srand(time(NULL));
	
	if(rank == 0){
		// generate vector
		for(i = 0; i < N; i++)
			vector[N] = rand()%1000;
	}
	
	MPI_Bcast(vector, N, MPI_INT, 0, MPI_COMM_WORLD);
	
	if(rank == 0){
		int a[N*N], buf, finish, result[N];
		
		// generate data
		for(i = 0; i < N; i++){
			for(j = 0; j < N; j++){
				a[i*N+j] = rand()%1000;
			}
		}
		
		// send data
		if(size > N){
			printf("Too many procs\n");
			for(i = 0; i < N; i++){
				MPI_Send(&a[i*N], N, MPI_INT, i+1,i, MPI_COMM_WORLD);
			}
			while(i < size-1){
				printf("i = %d\n", i);
				MPI_Send(&buf, 1, MPI_INT, i+1, FINISH_TAG, MPI_COMM_WORLD);
				i++;
			}
			// receive
			finish = 0;
			while(finish < N){
				printf("recv \n");
				MPI_Recv(&buf, 1, MPI_INT, MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
				result[status.MPI_TAG] = buf; 
				MPI_Send(&buf, 1, MPI_INT, status.MPI_SOURCE, FINISH_TAG, MPI_COMM_WORLD);
				finish ++;
			}
		}
		
		
		else{
			// send initial row
			for(i = 0; i < size-1; i++){
				MPI_Send(&a[i*N], N, MPI_INT, i+1, i, MPI_COMM_WORLD);
			}
			
			// send remain row
			finish = 0;
			while(finish < size-1){
				// receive result
				MPI_Recv(&buf, 1, MPI_INT, MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
				result[status.MPI_TAG] = buf; 
				printf("Recv from %d row %d\n", status.MPI_SOURCE, status.MPI_TAG);
				if(i < N){
					// send row
					MPI_Send(&a[i*N], N, MPI_INT, status.MPI_SOURCE, i, MPI_COMM_WORLD);		
				} else{
					printf("send finish %d\n", status.MPI_SOURCE);
					// send finish
					MPI_Send(&buf, 1, MPI_INT, status.MPI_SOURCE, FINISH_TAG, MPI_COMM_WORLD);
					finish ++;
				}
				i++;
			}
		}
		for(i = 0; i< N; i++)
			printf("result[%d] = %d\n", i, result[i]);
	}	
	else{
		int res = 0;
		while(1){
			// recv row
			MPI_Recv(&row, N, MPI_INT, 0, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
			//printf("rank %d recv row %d\n",rank, status.MPI_TAG);
			if(status.MPI_TAG!= FINISH_TAG){
				// cal result
				for(i = 0; i < N; i++)
					res = row[i]*vector[i];
				// send result
				MPI_Send(&res, 1, MPI_INT, 0, status.MPI_TAG, MPI_COMM_WORLD);
			}else
				break;
		}
	}
		

	MPI_Finalize();
	return 0;
}
